# -*- coding: utf-8 -*-  
'''
训练seq2seq

@author: luoyi
Created on 2021年3月18日
'''
import sys
import os
#    取项目根目录
ROOT_PATH = os.path.abspath(os.path.dirname(__file__)).split('nlp_seq2seq')[0]
ROOT_PATH = ROOT_PATH + "nlp_seq2seq"
sys.path.append(ROOT_PATH)

import data.dataset_couplet as ds_couplet
import utils.conf as conf
import utils.dictionaries as dictionaries
from models.seq2seq import Seq2SeqModel
from math import ceil


#    初始化字典
dictionaries.load_dict_from_file()

#    加载模型
model = Seq2SeqModel(learning_rate=conf.MODELS.get_learning_rate(),
                     worddict_size=dictionaries.dict_size(),
                     embedding_dims=conf.MODELS.get_embedding_dims(),
                     hidden_units=conf.MODELS.get_hidden_units())
model.show_info()

#    加载训练集、验证集
batch_size = conf.DATASET.get_batch_size()
count = conf.DATASET.get_count_train()
epochs = conf.DATASET.get_epochs()
shuffle_buffer_rate = conf.DATASET.get_shuffle_buffer_rate()
steps_per_epoch = ceil(count / batch_size)

db_train = ds_couplet.tensor_db(in_path=conf.DATASET.get_in_train(), 
                                out_path=conf.DATASET.get_label_train(), 
                                count=count, 
                                max_len=conf.TEXT.get_max_len(), 
                                batch_size=batch_size,
                                epochs=epochs,
                                shuffle_buffer_rate=shuffle_buffer_rate,
                                x_preprocess=None, 
                                y_preprocess=None)

db_val = ds_couplet.tensor_db(in_path=conf.DATASET.get_in_val(), 
                              out_path=conf.DATASET.get_label_val(), 
                              count=conf.DATASET.get_count_val(), 
                              max_len=conf.TEXT.get_max_len(), 
                              batch_size=batch_size,
                              epochs=epochs,
                              shuffle_buffer_rate=shuffle_buffer_rate,
                              x_preprocess=None, 
                              y_preprocess=None)

#    喂数据
model.train_tensor_db(db_train, db_val, 
                      steps_per_epoch, 
                      batch_size=batch_size, 
                      epochs=conf.DATASET.get_epochs(), 
                      auto_save_weights_after_traind=True, 
                      auto_save_weights_dir=conf.MODELS.get_model_save_weights(), 
                      auto_learning_rate_schedule=True, 
                      auto_tensorboard=True, 
                      auto_tensorboard_dir=conf.MODELS.get_tensorboard_dir())